"""
Milvus 向量数据库客户端

兼容层，用于提供 MilvusClient 接口，实际可接入 Milvus 或 Zilliz。
"""
from typing import List, Optional, Dict, Any, Union, Tuple
import numpy as np
import json
import logging
import time
from utils.logger import Logger
from utils.constants import MILVUS_HOST, MILVUS_PORT, MILVUS_DB_NAME, MILVUS_VECTOR_DIM
from middleware.retry_utils import retry, async_retry, vector_db_retry, VECTOR_DB_RETRY_CONFIG, task_scheduler
from middleware.data_consistency import data_consistency_manager, transaction

logger = Logger.get_logger("milvus_client")

# 尝试导入 Milvus SDK，失败时提供优雅降级
pymilvus_available = False
try:
    from pymilvus import (
        connections,
        Collection,
        utility,
        FieldSchema,
        CollectionSchema,
        DataType
    )
    from pymilvus import db as milvus_db
    pymilvus_available = True
except ImportError:
    logger.warning("Milvus SDK (pymilvus) 未安装，将使用降级模式")
    
    # 创建兼容接口
    class MockConnections:
        def connect(self, *args, **kwargs): pass
        def disconnect(self, *args, **kwargs): pass
        def get_connection(self, *args, **kwargs): return None
    
    class MockUtility:
        def has_collection(self, *args, **kwargs): return False
        def drop_collection(self, *args, **kwargs): pass
    
    class MockCollection:
        def __init__(self, *args, **kwargs): pass
        def create_index(self, *args, **kwargs): pass
        def load(self, *args, **kwargs): pass
        def insert(self, *args, **kwargs): return None
        def search(self, *args, **kwargs): return []
        def query(self, *args, **kwargs): return []
        def flush(self, *args, **kwargs): pass
        def release(self, *args, **kwargs): pass
    
    class MockDB:
        def create_database(self, *args, **kwargs): pass
        def use_database(self, *args, **kwargs): pass
    
    connections = MockConnections()
    utility = MockUtility()
    Collection = MockCollection
    milvus_db = MockDB()


class MilvusClient:
    """完整实现的 Milvus 客户端封装"""

    def __init__(self, uri: str = None, host: str = None, port: int = None, db_name: str = None, **kwargs):
        """
        初始化 Milvus 客户端
        :param uri: Milvus 服务地址，默认 None
        :param host: Milvus 主机地址，默认使用配置中的 MILVUS_HOST
        :param port: Milvus 端口，默认使用配置中的 MILVUS_PORT
        :param db_name: 数据库名称，默认使用配置中的 MILVUS_DB_NAME
        :param kwargs: 其他连接参数
        """
        self.host = host or MILVUS_HOST
        self.port = port or MILVUS_PORT
        self.db_name = db_name or MILVUS_DB_NAME
        self.uri = uri or f"http://{self.host}:{self.port}"
        self._connected = False
        self._collections: Dict[str, Collection] = {}  # 缓存已加载的集合
        self._connect(**kwargs)
        
    @vector_db_retry
    def _connect(self, **kwargs):
        """建立连接并初始化（带重试机制）"""
        try:
            logger.info(f"正在连接 Milvus 服务: {self.host}:{self.port}, 数据库: {self.db_name}")
            
            # 连接 Milvus 服务
            connections.connect(
                alias="default",
                host=self.host,
                port=self.port,
                **kwargs
            )
            
            # 使用指定数据库
            if self.db_name != "default":
                try:
                    milvus_db.use_database(self.db_name)
                except Exception:
                    # 数据库不存在时创建
                    logger.info(f"数据库 {self.db_name} 不存在，尝试创建")
                    milvus_db.create_database(self.db_name)
                    milvus_db.use_database(self.db_name)
            
            self._connected = True
            logger.info("Milvus 连接成功")
            
        except Exception as e:
            self._connected = False
            logger.error(f"Milvus 连接失败: {str(e)}")
            raise  # 抛出异常以便重试装饰器捕获并重试

    @vector_db_retry
    def search(
        self,
        collection_name: str,
        query_vectors: List[List[float]],
        top_k: int = 10,
        metric_type: str = "COSINE",
        output_fields: Optional[List[str]] = None,
        filter_expr: Optional[str] = None,
        **kwargs
    ) -> List[List["SearchResult"]]:
        """
        向量检索（带重试机制）
        :param collection_name: 集合名称
        :param query_vectors: 查询向量列表
        :param top_k: 返回 topK 结果
        :param metric_type: 度量方式，默认 COSINE
        :param output_fields: 需要返回的字段列表
        :param filter_expr: 过滤表达式
        :return: 检索结果，格式兼容官方 SDK
        """
        if not self._connected or not pymilvus_available:
            logger.warning("Milvus 未连接或 SDK 不可用，返回空结果")
            return [[SearchResult(id=i, score=0.0)] for i in range(top_k)]
        
        try:
            # 获取或加载集合
            collection = self._get_or_load_collection(collection_name)
            if not collection:
                logger.warning(f"集合 {collection_name} 不存在或无法加载")
                return [[SearchResult(id=i, score=0.0)] for i in range(top_k)]
            
            # 执行搜索
            search_params = {"metric_type": metric_type, "params": {"nprobe": 10}}
            results = collection.search(
                data=query_vectors,
                anns_field="embedding",
                param=search_params,
                limit=top_k,
                expr=filter_expr,
                output_fields=output_fields or [],
                consistency_level="Strong",
                **kwargs
            )
            
            # 转换结果格式
            search_results = []
            for hits in results:
                hit_list = []
                for hit in hits:
                    result = SearchResult(
                        id=hit.id,
                        score=hit.score,
                        entity=hit.entity.to_dict() if hasattr(hit, 'entity') else {}
                    )
                    hit_list.append(result)
                search_results.append(hit_list)
            
            logger.debug(f"Milvus 搜索完成: 集合={collection_name}, 查询向量数={len(query_vectors)}, top_k={top_k}")
            return search_results
            
        except Exception as e:
            logger.error(f"Milvus 搜索失败: {str(e)}")
            raise  # 抛出异常以便重试装饰器捕获并重试

    @vector_db_retry
    def get_entity_by_id(
        self,
        collection_name: str,
        ids: Union[List[int], List[str]],
        output_fields: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """
        根据 ID 获取实体（带重试机制）
        :param collection_name: 集合名称
        :param ids: 实体 ID 列表
        :param output_fields: 返回字段，默认返回全部
        :return: 实体字典列表
        """
        if not self._connected or not pymilvus_available:
            logger.warning("Milvus 未连接或 SDK 不可用，返回空结果")
            return [{} for _ in ids]
        
        try:
            # 对于大量ID查询，分批处理
            if len(ids) > 100:
                logger.info(f"查询ID数量较多 ({len(ids)}), 将分批处理")
                batch_size = 100
                ordered_results = []
                
                # 先创建一个空结果列表，保持与输入ID相同顺序
                temp_results = [{} for _ in ids]
                # 创建ID到索引的映射
                id_to_index = {str(id_val): idx for idx, id_val in enumerate(ids)}
                
                for i in range(0, len(ids), batch_size):
                    batch_ids = ids[i:i+batch_size]
                    batch_results = self._get_entities_by_id_batch(collection_name, batch_ids, output_fields)
                    
                    # 将批次结果放入对应位置
                    for idx, result in enumerate(batch_results):
                        id_val = str(batch_ids[idx])
                        if id_val in id_to_index:
                            temp_results[id_to_index[id_val]] = result
                
                return temp_results
            else:
                return self._get_entities_by_id_batch(collection_name, ids, output_fields)
                
        except Exception as e:
            logger.error(f"Milvus ID查询失败: {str(e)}")
            raise  # 抛出异常以便重试装饰器捕获并重试
    
    def _get_entities_by_id_batch(
        self,
        collection_name: str,
        ids: Union[List[int], List[str]],
        output_fields: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """
        内部方法：批量查询实体
        :param collection_name: 集合名称
        :param ids: 实体 ID 列表
        :param output_fields: 返回字段，默认返回全部
        :return: 实体字典列表
        """
        # 获取或加载集合
        collection = self._get_or_load_collection(collection_name)
        if not collection:
            return [{} for _ in ids]
        
        # 构建查询表达式
        if all(isinstance(id, int) for id in ids):
            id_str = ", ".join(map(str, ids))
        else:
            id_str = ", ".join(f'"{id}"' for id in ids)
        
        expr = f"id in [{id_str}]"
        
        # 执行查询
        results = collection.query(
            expr=expr,
            output_fields=output_fields or [],
            consistency_level="Strong"
        )
        
        # 构建 ID 到结果的映射，确保返回顺序与请求顺序一致
        result_map = {str(r.get("id")): r for r in results}
        ordered_results = []
        
        for id_val in ids:
            ordered_results.append(result_map.get(str(id_val), {}))
        
        return ordered_results

    @vector_db_retry
    def insert(self, collection_name: str, data: List[Dict[str, Any]], **kwargs) -> Any:
        """
        插入数据（带重试机制）
        :param collection_name: 集合名称
        :param data: 数据列表
        :return: 插入结果
        """
        if not self._connected or not pymilvus_available:
            logger.warning("Milvus 未连接或 SDK 不可用，跳过插入")
            return None
        
        try:
            # 获取或加载集合
            collection = self._get_or_load_collection(collection_name)
            if not collection:
                logger.error(f"集合 {collection_name} 不存在或无法加载")
                return None
            
            # 批量插入，每次最多插入1000条
            batch_size = kwargs.get("batch_size", 1000)
            all_ids = []
            
            for i in range(0, len(data), batch_size):
                batch = data[i:i + batch_size]
                result = collection.insert(batch)
                if result and hasattr(result, 'primary_keys'):
                    all_ids.extend(result.primary_keys)
                
                # 每批次插入后刷新以确保数据可见
                collection.flush()
            
            logger.info(f"Milvus 插入完成: 集合={collection_name}, 总插入数={len(data)}, 批次={len(data) // batch_size + 1}")
            return {"insert_count": len(data), "primary_keys": all_ids}
            
        except Exception as e:
            logger.error(f"Milvus 插入失败: {str(e)}")
            raise  # 抛出异常以便重试装饰器捕获并重试

    def create_collection(self, collection_name: str, schema: Optional[Dict[str, Any]] = None, **kwargs) -> bool:
        """
        创建集合
        :param collection_name: 集合名称
        :param schema: 集合模式，包含字段定义
        :return: 是否成功
        """
        if not self._connected or not pymilvus_available:
            logger.warning("Milvus 未连接或 SDK 不可用，跳过创建集合")
            return False
        
        try:
            # 检查集合是否已存在
            if utility.has_collection(collection_name):
                logger.info(f"集合 {collection_name} 已存在")
                # 获取并返回现有集合
                self._collections[collection_name] = Collection(collection_name)
                return True
            
            # 如果没有提供schema，使用默认schema
            if schema is None:
                # 默认的向量集合schema
                fields = [
                    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
                    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=kwargs.get("dim", MILVUS_VECTOR_DIM)),
                    FieldSchema(name="metadata", dtype=DataType.JSON)
                ]
                schema_obj = CollectionSchema(fields, description=f"Vector collection {collection_name}")
            else:
                # 根据提供的schema构建字段
                fields = []
                for field_name, field_info in schema.items():
                    field_type_map = {
                        "int": DataType.INT64,
                        "str": DataType.VARCHAR,
                        "float": DataType.FLOAT,
                        "vector": DataType.FLOAT_VECTOR,
                        "json": DataType.JSON
                    }
                    
                    dtype = field_type_map.get(field_info.get("type"), DataType.VARCHAR)
                    kwargs = {
                        "name": field_name,
                        "dtype": dtype,
                        "is_primary": field_info.get("primary", False),
                        "auto_id": field_info.get("auto_id", False)
                    }
                    
                    if dtype == DataType.VARCHAR:
                        kwargs["max_length"] = field_info.get("max_length", 65535)
                    elif dtype == DataType.FLOAT_VECTOR:
                        kwargs["dim"] = field_info.get("dim", MILVUS_VECTOR_DIM)
                    
                    fields.append(FieldSchema(**kwargs))
                
                schema_obj = CollectionSchema(fields, description=schema.get("description", f"Collection {collection_name}"))
            
            # 创建集合
            collection = Collection(name=collection_name, schema=schema_obj)
            logger.info(f"Milvus 集合创建成功: {collection_name}")
            
            # 创建向量索引
            if any(f.name == "embedding" and f.dtype == DataType.FLOAT_VECTOR for f in schema_obj.fields):
                index_params = {
                    "index_type": "IVF_FLAT",
                    "metric_type": kwargs.get("metric_type", "COSINE"),
                    "params": {"nlist": 128}
                }
                collection.create_index("embedding", index_params)
                logger.info(f"Milvus 向量索引创建成功: {collection_name}")
            
            # 缓存集合对象
            self._collections[collection_name] = collection
            return True
            
        except Exception as e:
            logger.error(f"Milvus 集合创建失败: {str(e)}")
            return False
    
    def _get_or_load_collection(self, collection_name: str) -> Optional[Collection]:
        """
        获取或加载集合
        :param collection_name: 集合名称
        :return: 集合对象
        """
        # 检查缓存中是否已有
        if collection_name in self._collections:
            return self._collections[collection_name]
        
        # 检查集合是否存在
        if not utility.has_collection(collection_name):
            logger.error(f"集合 {collection_name} 不存在")
            return None
        
        # 加载集合
        try:
            collection = Collection(collection_name)
            collection.load()  # 加载到内存
            self._collections[collection_name] = collection
            logger.info(f"Milvus 集合加载成功: {collection_name}")
            return collection
        except Exception as e:
            logger.error(f"Milvus 集合加载失败: {str(e)}")
            return None
    
    def drop_collection(self, collection_name: str) -> bool:
        """
        删除集合
        :param collection_name: 集合名称
        :return: 是否成功
        """
        if not self._connected or not pymilvus_available:
            logger.warning("Milvus 未连接或 SDK 不可用，跳过删除")
            return False
        
        try:
            # 从缓存中移除
            if collection_name in self._collections:
                del self._collections[collection_name]
            
            # 删除集合
            if utility.has_collection(collection_name):
                utility.drop_collection(collection_name)
                logger.info(f"Milvus 集合删除成功: {collection_name}")
                return True
            else:
                logger.info(f"集合 {collection_name} 不存在")
                return True
        except Exception as e:
            logger.error(f"Milvus 集合删除失败: {str(e)}")
            return False
    
    def has_collection(self, collection_name: str) -> bool:
        """
        检查集合是否存在
        :param collection_name: 集合名称
        :return: 是否存在
        """
        if not self._connected or not pymilvus_available:
            return False
        
        try:
            return utility.has_collection(collection_name)
        except Exception as e:
            logger.error(f"检查集合失败: {str(e)}")
            return False
    
    def get_collection_stats(self, collection_name: str) -> Optional[Dict[str, Any]]:
        """
        获取集合统计信息
        :param collection_name: 集合名称
        :return: 统计信息
        """
        if not self._connected or not pymilvus_available:
            return None
        
        try:
            collection = self._get_or_load_collection(collection_name)
            if not collection:
                return None
            
            stats = {
                "collection_name": collection_name,
                "row_count": collection.num_entities,
                "is_loaded": collection.is_loaded
            }
            return stats
        except Exception as e:
            logger.error(f"获取集合统计失败: {str(e)}")
            return None
    
    def close(self):
        """
        关闭连接并释放资源
        """
        try:
            # 释放所有已加载的集合
            for collection_name, collection in self._collections.items():
                try:
                    if collection.is_loaded:
                        collection.release()
                except Exception:
                    pass
            
            # 清空集合缓存
            self._collections.clear()
            
            # 断开连接
            if self._connected:
                connections.disconnect("default")
                self._connected = False
                logger.info("Milvus 连接已关闭")
        except Exception as e:
            logger.error(f"关闭 Milvus 连接失败: {str(e)}")


class SearchResult:
    """检索结果项"""
    def __init__(self, id: Union[int, str], score: float, entity: Optional[Dict[str, Any]] = None):
        self.id = id
        self.score = score
        self.entity = entity or {}

    def __repr__(self):
        return f"SearchResult(id={self.id}, score={self.score})"
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "id": self.id,
            "score": self.score,
            "entity": self.entity
        }


# 全局 Milvus 客户端实例
_milvus_client_instance = None

def get_milvus_client() -> MilvusClient:
    """
    获取全局 Milvus 客户端实例
    :return: MilvusClient 实例
    """
    global _milvus_client_instance
    if _milvus_client_instance is None:
        _milvus_client_instance = MilvusClient()
    return _milvus_client_instance


# 熔断器实现
class CircuitBreaker:
    """
    熔断器模式实现
    支持三种状态：关闭(closed)、打开(open)、半开(half-open)
    """
    def __init__(self, failure_threshold: int = 5, reset_timeout: float = 30.0, 
                 recovery_threshold: int = 2, name: str = "default"):
        """
        初始化熔断器
        :param failure_threshold: 失败阈值，超过此值进入打开状态
        :param reset_timeout: 重置超时时间（秒），打开状态持续时间
        :param recovery_threshold: 恢复阈值，半开状态下成功次数达到此值恢复到关闭状态
        :param name: 熔断器名称，用于日志
        """
        self.failure_threshold = failure_threshold
        self.reset_timeout = reset_timeout
        self.recovery_threshold = recovery_threshold
        self.name = name
        self.state = "closed"  # closed, open, half-open
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time = None
        
    def is_allowed(self) -> bool:
        """
        判断是否允许请求通过
        :return: 是否允许
        """
        current_time = time.time()
        
        if self.state == "open":
            # 检查是否可以进入半开状态
            if self.last_failure_time and current_time - self.last_failure_time >= self.reset_timeout:
                logger.info(f"熔断器 {self.name} 从打开状态进入半开状态")
                self.state = "half-open"
                self.success_count = 0
            else:
                logger.warning(f"熔断器 {self.name} 处于打开状态，拒绝请求")
                return False
        
        return True
    
    def record_success(self):
        """
        记录请求成功
        """
        if self.state == "closed":
            # 关闭状态下，重置失败计数
            self.failure_count = 0
        elif self.state == "half-open":
            # 半开状态下，增加成功计数
            self.success_count += 1
            if self.success_count >= self.recovery_threshold:
                logger.info(f"熔断器 {self.name} 从半开状态恢复到关闭状态")
                self.state = "closed"
                self.failure_count = 0
                self.success_count = 0
    
    def record_failure(self):
        """
        记录请求失败
        """
        self.failure_count += 1
        self.last_failure_time = time.time()
        
        if self.state == "closed":
            # 关闭状态下，检查是否需要打开
            if self.failure_count >= self.failure_threshold:
                logger.warning(f"熔断器 {self.name} 从关闭状态进入打开状态，失败次数: {self.failure_count}")
                self.state = "open"
        elif self.state == "half-open":
            # 半开状态下失败，直接回到打开状态
            logger.warning(f"熔断器 {self.name} 从半开状态回到打开状态")
            self.state = "open"
            self.failure_count = 0
    
    def get_state(self) -> Dict[str, Any]:
        """
        获取熔断器状态信息
        :return: 状态信息字典
        """
        return {
            "name": self.name,
            "state": self.state,
            "failure_count": self.failure_count,
            "success_count": self.success_count,
            "last_failure_time": self.last_failure_time
        }

# 数据路由管理（Milvus 与 PostgreSQL+PgVector 的自动切换）
class VectorDatabaseRouter:
    """
    向量数据库路由器，根据数据规模自动选择 Milvus 或 PostgreSQL+PgVector
    集成熔断降级策略、重试机制和任务调度
    """
    
    def __init__(self, small_threshold: int = 1000):
        """
        初始化路由器
        :param small_threshold: 小型数据集阈值（低于此值使用 PostgreSQL）
        """
        self.small_threshold = small_threshold
        self.milvus_client = get_milvus_client()
        self._collection_sizes = {}
        
        # 初始化熔断器
        self.circuit_breakers = {
            "milvus_search": CircuitBreaker(failure_threshold=3, reset_timeout=60.0, name="milvus_search"),
            "milvus_insert": CircuitBreaker(failure_threshold=3, reset_timeout=60.0, name="milvus_insert"),
            "milvus_query": CircuitBreaker(failure_threshold=3, reset_timeout=60.0, name="milvus_query")
        }
        
        # 导入PostgreSQL向量查询相关模块
        self.postgresql_available = False
        try:
            # 这里需要导入实际的PostgreSQL向量查询模块
            # 暂时使用占位符，实际实现时需要替换
            self.postgresql_available = True
        except ImportError:
            logger.warning("PostgreSQL向量查询模块未安装，降级功能受限")
        
        # 初始化任务调度器
        self._initialize_task_scheduler()
    
    @vector_db_retry
    def should_use_postgresql(self, collection_name: str, expected_size: Optional[int] = None) -> bool:
        """
        判断是否应该使用 PostgreSQL
        :param collection_name: 集合/表名
        :param expected_size: 预期数据量
        :return: 是否使用 PostgreSQL
        """
        # 如果 Milvus 不可用，降级到 PostgreSQL
        if not self.milvus_client._connected:
            logger.warning(f"Milvus 不可用，使用 PostgreSQL: {collection_name}")
            return True
        
        # 如果指定了预期大小且小于阈值，使用 PostgreSQL
        if expected_size is not None and expected_size < self.small_threshold:
            logger.info(f"小型数据集，使用 PostgreSQL: {collection_name}, 预期大小={expected_size}")
            return True
        
        # 检查已知的集合大小
        if collection_name in self._collection_sizes:
            if self._collection_sizes[collection_name] < self.small_threshold:
                return True
        
        # 默认使用 Milvus
        logger.info(f"大型数据集，使用 Milvus: {collection_name}")
        return False
    
    def update_collection_size(self, collection_name: str, size: int):
        """
        更新集合大小记录
        :param collection_name: 集合/表名
        :param size: 集合大小
        """
        self._collection_sizes[collection_name] = size
        logger.debug(f"更新集合大小: {collection_name} = {size}")
    
    @vector_db_retry
    def search(self, collection_name: str, query_vectors: List[List[float]], top_k: int = 10, 
               metric_type: str = "COSINE", filter: Optional[Dict] = None, 
               db_type: Optional[str] = None, vector_field: str = "embedding", **kwargs) -> List[List[SearchResult]]:
        """
        向量检索（支持熔断降级和自动重试）
        :param collection_name: 集合名称
        :param query_vectors: 查询向量列表
        :param top_k: 返回 topK 结果
        :param metric_type: 度量方式，默认 COSINE
        :param filter: 过滤条件
        :param db_type: 数据库类型，可选值: "milvus", "postgresql"
        :param vector_field: 向量字段名
        :return: 检索结果
        """
        try:
            # 确定使用的数据库类型
            if db_type is None:
                use_postgresql = self.should_use_postgresql(collection_name)
                db_type = "postgresql" if use_postgresql else "milvus"
            
            # 对于 Milvus，检查熔断器状态
            if db_type == "milvus":
                if not self.circuit_breakers["milvus_search"].is_allowed():
                    # 熔断器打开，降级到 PostgreSQL
                    logger.warning(f"Milvus搜索熔断器打开，降级到PostgreSQL: {collection_name}")
                    db_type = "postgresql"
            
            # 执行搜索
            if db_type == "milvus":
                # 使用 Milvus 进行搜索
                results = self._search_with_milvus(collection_name, query_vectors, top_k, metric_type, filter, vector_field, **kwargs)
                self.circuit_breakers["milvus_search"].record_success()
                return results
            else:
                # 使用 PostgreSQL 进行搜索（如果可用）
                if self.postgresql_available:
                    return self._search_with_postgresql(collection_name, query_vectors, top_k, metric_type, filter, **kwargs)
                else:
                    # PostgreSQL 不可用时的最终降级策略
                    logger.warning(f"PostgreSQL不可用，返回空结果: {collection_name}")
                    return [[SearchResult(id=str(i), score=0.0)] for _ in query_vectors for i in range(top_k)]
                    
        except Exception as e:
            # 记录错误并返回空结果
            logger.error(f"向量搜索失败: {str(e)}")
            if db_type == "milvus":
                self.circuit_breakers["milvus_search"].record_failure()
            # 抛出异常以便重试装饰器捕获
            raise
    
    @vector_db_retry
    def get_entities_by_id(self, collection_name: str, ids: List[Union[int, str]], 
                          output_fields: Optional[List[str]] = None, db_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        根据ID获取实体（支持熔断降级和自动重试）
        :param collection_name: 集合名称
        :param ids: 实体ID列表
        :param output_fields: 返回字段列表
        :param db_type: 数据库类型，可选值: "milvus", "postgresql"
        :return: 实体字典列表
        """
        try:
            # 确定使用的数据库类型
            if db_type is None:
                use_postgresql = self.should_use_postgresql(collection_name)
                db_type = "postgresql" if use_postgresql else "milvus"
            
            # 对于 Milvus，检查熔断器状态
            if db_type == "milvus":
                if not self.circuit_breakers["milvus_query"].is_allowed():
                    # 熔断器打开，降级到 PostgreSQL
                    logger.warning(f"Milvus查询熔断器打开，降级到PostgreSQL: {collection_name}")
                    db_type = "postgresql"
            
            # 执行查询
            if db_type == "milvus":
                # 使用 Milvus 进行查询
                results = self.milvus_client.get_entity_by_id(collection_name, ids, output_fields)
                self.circuit_breakers["milvus_query"].record_success()
                return results
            else:
                # 使用 PostgreSQL 进行查询（如果可用）
                if self.postgresql_available:
                    return self._get_entities_with_postgresql(collection_name, ids, output_fields)
                else:
                    # PostgreSQL 不可用时的最终降级策略
                    logger.warning(f"PostgreSQL不可用，返回空结果: {collection_name}")
                    return [{} for _ in ids]
                    
        except Exception as e:
            # 记录错误并返回空结果
            logger.error(f"获取实体失败: {str(e)}")
            if db_type == "milvus":
                self.circuit_breakers["milvus_query"].record_failure()
            # 抛出异常以便重试装饰器捕获
            raise
    
    @vector_db_retry
    def insert(self, collection_name: str, data: List[Dict[str, Any]], 
               db_type: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        插入数据（支持熔断降级、自动重试和事务一致性）
        :param collection_name: 集合名称
        :param data: 数据列表
        :param db_type: 数据库类型，可选值: "milvus", "postgresql"
        :return: 插入结果
        """
        # 生成事务ID
        transaction_id = kwargs.pop("transaction_id", None)
        
        # 如果提供了事务ID，使用事务一致性管理器
        if transaction_id:
            try:
                # 添加操作到事务
                data_consistency_manager.add_operation(transaction_id, {
                    "type": "vector_insert",
                    "collection_name": collection_name,
                    "data": data,
                    "db_type": db_type,
                    "kwargs": kwargs
                })
                
                # 先执行操作，但不提交（实际提交由事务管理器处理）
                result = self._insert_without_transaction(collection_name, data, db_type, **kwargs)
                
                # 记录操作结果用于事务恢复
                data_consistency_manager.add_operation(transaction_id, {
                    "type": "vector_insert_result",
                    "collection_name": collection_name,
                    "result": result,
                    "db_type": db_type
                })
                
                return result
                
            except Exception as e:
                # 记录错误并抛出，供事务管理器处理回滚
                logger.error(f"事务内数据插入失败: {str(e)}")
                raise
        else:
            # 无事务的普通插入
            return self._insert_without_transaction(collection_name, data, db_type, **kwargs)
    
    def _insert_without_transaction(self, collection_name: str, data: List[Dict[str, Any]],
                                  db_type: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        无事务插入操作的实际实现
        """
        try:
            # 确定使用的数据库类型
            if db_type is None:
                use_postgresql = self.should_use_postgresql(collection_name, expected_size=len(data))
                db_type = "postgresql" if use_postgresql else "milvus"
            
            # 对于 Milvus，检查熔断器状态
            if db_type == "milvus":
                if not self.circuit_breakers["milvus_insert"].is_allowed():
                    # 熔断器打开，降级到 PostgreSQL
                    logger.warning(f"Milvus插入熔断器打开，降级到PostgreSQL: {collection_name}")
                    db_type = "postgresql"
            
            # 执行插入
            if db_type == "milvus":
                # 使用 Milvus 进行插入
                result = self.milvus_client.insert(collection_name, data, **kwargs)
                self.circuit_breakers["milvus_insert"].record_success()
                
                # 更新集合大小
                if result and result.get("insert_count"):
                    self.update_collection_size(collection_name, 
                                              self._collection_sizes.get(collection_name, 0) + result["insert_count"])
                
                return result
            else:
                # 使用 PostgreSQL 进行插入（如果可用）
                if self.postgresql_available:
                    result = self._insert_with_postgresql(collection_name, data, **kwargs)
                    
                    # 更新集合大小
                    if result and result.get("insert_count"):
                        self.update_collection_size(collection_name, 
                                                  self._collection_sizes.get(collection_name, 0) + result["insert_count"])
                    
                    return result
                else:
                    # PostgreSQL 不可用时的最终降级策略
                    logger.warning(f"PostgreSQL不可用，跳过插入: {collection_name}")
                    return {"insert_count": 0, "error": "存储不可用"}
                    
        except Exception as e:
            # 记录错误并返回失败结果
            logger.error(f"数据插入失败: {str(e)}")
            if db_type == "milvus":
                self.circuit_breakers["milvus_insert"].record_failure()
            # 抛出异常以便重试装饰器捕获
            raise
    
    def sync_collections(self, source_db: str, target_db: str, 
                        collection_name: str, sync_type: str = "full") -> Dict[str, Any]:
        """
        同步两个数据库之间的集合数据
        :param source_db: 源数据库类型
        :param target_db: 目标数据库类型
        :param collection_name: 集合名称
        :param sync_type: 同步类型: "full" 或 "incremental"
        :return: 同步结果
        """
        logger.info(f"开始同步集合数据: {source_db} -> {target_db}, 集合: {collection_name}")
        
        try:
            # 获取源数据库的所有数据
            if source_db == "milvus":
                # 从Milvus获取数据
                source_data = self._get_all_milvus_data(collection_name)
            else:
                # 从PostgreSQL获取数据
                source_data = self._get_all_postgresql_data(collection_name)
            
            # 使用数据一致性管理器执行同步
            sync_result = data_consistency_manager.sync_data(
                source=source_db,
                target=target_db,
                data=source_data,
                sync_type=sync_type
            )
            
            return sync_result
            
        except Exception as e:
            logger.error(f"集合同步失败: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def _get_all_milvus_data(self, collection_name: str) -> List[Dict[str, Any]]:
        """
        获取Milvus集合中的所有数据
        """
        # 实际实现需要根据Milvus客户端的API进行调整
        logger.info(f"从Milvus获取集合 {collection_name} 的所有数据")
        # 这里应该实现查询Milvus集合所有数据的逻辑
        return []
    
    def _get_all_postgresql_data(self, collection_name: str) -> List[Dict[str, Any]]:
        """
        获取PostgreSQL表中的所有数据
        """
        # 实际实现需要根据PostgreSQL客户端的API进行调整
        logger.info(f"从PostgreSQL获取表 {collection_name} 的所有数据")
        # 这里应该实现查询PostgreSQL表所有数据的逻辑
        return []
    
    def verify_data_consistency(self, collection_name: str) -> Dict[str, Any]:
        """
        验证不同数据库之间的数据一致性
        :param collection_name: 集合/表名称
        :return: 一致性检查结果
        """
        logger.info(f"开始验证数据一致性: {collection_name}")
        
        try:
            # 从两个数据库获取数据
            milvus_data = []
            postgres_data = []
            
            # 检查Milvus数据
            if self.milvus_client._connected:
                milvus_data = self._get_all_milvus_data(collection_name)
            
            # 检查PostgreSQL数据
            if self.postgresql_available:
                postgres_data = self._get_all_postgresql_data(collection_name)
            
            # 比较数据一致性
            consistency_result = self._compare_data_consistency(
                milvus_data, postgres_data,
                "milvus", "postgresql"
            )
            
            # 如果发现不一致，记录并尝试恢复
            if not consistency_result["consistent"]:
                logger.warning(f"数据不一致: {collection_name}, 差异: {consistency_result['differences']}")
                
                # 触发自动恢复
                recovery_result = data_consistency_manager.recover_from_failure({
                    "failure_type": "data_inconsistency",
                    "collection_name": collection_name,
                    "differences": consistency_result["differences"]
                })
                
                consistency_result["recovery_result"] = recovery_result
            
            return consistency_result
            
        except Exception as e:
            logger.error(f"数据一致性验证失败: {str(e)}")
            return {
                "consistent": False,
                "error": str(e)
            }
    
    def _compare_data_consistency(self, data1: List[Dict[str, Any]], data2: List[Dict[str, Any]],
                                source1: str, source2: str) -> Dict[str, Any]:
        """
        比较两组数据的一致性
        """
        # 这里应该实现数据比较逻辑
        # 例如：比较记录数量、比较每个记录的内容等
        differences = []
        
        # 简单比较记录数量
        if len(data1) != len(data2):
            differences.append({
                "type": "count_mismatch",
                f"{source1}_count": len(data1),
                f"{source2}_count": len(data2)
            })
        
        return {
            "consistent": len(differences) == 0,
            "differences": differences
        }
    
    @vector_db_retry
    def create_collection(self, collection_name: str, schema: Optional[Dict[str, Any]] = None, 
                         db_type: Optional[str] = None, **kwargs) -> bool:
        """
        创建集合（支持熔断降级和自动重试）
        :param collection_name: 集合名称
        :param schema: 集合模式
        :param db_type: 数据库类型，可选值: "milvus", "postgresql"
        :return: 是否成功
        """
        try:
            # 确定使用的数据库类型
            if db_type is None:
                use_postgresql = self.should_use_postgresql(collection_name)
                db_type = "postgresql" if use_postgresql else "milvus"
            
            # 执行创建
            if db_type == "milvus":
                # 使用 Milvus 创建集合
                return self.milvus_client.create_collection(collection_name, schema, **kwargs)
            else:
                # 使用 PostgreSQL 创建集合（如果可用）
                if self.postgresql_available:
                    return self._create_collection_with_postgresql(collection_name, schema, **kwargs)
                else:
                    logger.warning(f"PostgreSQL不可用，无法创建集合: {collection_name}")
                    return False
                    
        except Exception as e:
            logger.error(f"创建集合失败: {str(e)}")
            # 抛出异常以便重试装饰器捕获
            raise
    
    def _search_with_milvus(self, collection_name: str, query_vectors: List[List[float]], 
                           top_k: int = 10, metric_type: str = "COSINE", 
                           filter: Optional[Dict] = None, vector_field: str = "embedding", **kwargs) -> List[List[SearchResult]]:
        """
        使用 Milvus 执行向量搜索
        """
        # 转换过滤条件
        filter_expr = None
        if filter:
            # 简单的过滤条件转换，实际实现可能需要更复杂的逻辑
            conditions = []
            for field, value in filter.items():
                if isinstance(value, dict):
                    # 处理嵌套条件，如 {"$regex": "^doc_123_"}
                    if "$regex" in value:
                        conditions.append(f"{field} like '{value['$regex'].replace('^', '').replace('$', '%')}'")
                else:
                    conditions.append(f"{field} = '{value}'")
            filter_expr = " and ".join(conditions)
        
        # 调用 Milvus 客户端执行搜索
        return self.milvus_client.search(
            collection_name=collection_name,
            query_vectors=query_vectors,
            top_k=top_k,
            metric_type=metric_type,
            filter_expr=filter_expr,
            **kwargs
        )
    
    def _search_with_postgresql(self, collection_name: str, query_vectors: List[List[float]], 
                               top_k: int = 10, metric_type: str = "COSINE", 
                               filter: Optional[Dict] = None, **kwargs) -> List[List[SearchResult]]:
        """
        使用 PostgreSQL 执行向量搜索（占位实现）
        实际实现需要使用 pgvector 扩展
        """
        logger.info(f"使用PostgreSQL执行向量搜索: {collection_name}")
        # 占位实现，返回空结果
        return [[SearchResult(id=str(i), score=0.0)] for _ in query_vectors for i in range(top_k)]
    
    def _get_entities_with_postgresql(self, collection_name: str, ids: List[Union[int, str]], 
                                     output_fields: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """
        使用 PostgreSQL 根据ID获取实体（占位实现）
        """
        logger.info(f"使用PostgreSQL执行ID查询: {collection_name}")
        # 占位实现，返回空结果
        return [{} for _ in ids]
    
    def _insert_with_postgresql(self, collection_name: str, data: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
        """
        使用 PostgreSQL 插入数据（占位实现）
        """
        logger.info(f"使用PostgreSQL执行数据插入: {collection_name}, 数量={len(data)}")
        # 占位实现，返回成功结果
        return {"insert_count": len(data), "primary_keys": [i for i in range(len(data))]}
    
    def _create_collection_with_postgresql(self, collection_name: str, 
                                         schema: Optional[Dict[str, Any]] = None, **kwargs) -> bool:
        """
        使用 PostgreSQL 创建集合（占位实现）
        """
        logger.info(f"使用PostgreSQL创建集合: {collection_name}")
        # 占位实现，返回成功
        return True
    
    def get_circuit_breaker_state(self) -> Dict[str, Dict[str, Any]]:
        """
        获取所有熔断器的状态
        :return: 熔断器状态字典
        """
        return {
            name: breaker.get_state()
            for name, breaker in self.circuit_breakers.items()
        }
    
    def _initialize_task_scheduler(self):
        """
        初始化任务调度器
        添加健康检查和集合大小更新等定时任务
        """
        # 启动健康检查定时任务
        task_scheduler.add_interval_task(
            func=self._health_check,
            seconds=60,  # 每分钟执行一次健康检查
            task_id="milvus_health_check"
        )
        
        # 启动集合大小更新定时任务
        task_scheduler.add_interval_task(
            func=self._update_all_collection_sizes,
            seconds=300,  # 每5分钟更新一次集合大小
            task_id="update_collection_sizes"
        )
        
        logger.info("向量数据库任务调度器初始化完成")
    
    def _health_check(self):
        """
        执行健康检查任务
        检查 Milvus 和 PostgreSQL 的连接状态
        """
        try:
            # 检查 Milvus 连接状态
            if not self.milvus_client._connected:
                logger.warning("Milvus 连接断开，尝试重新连接")
                try:
                    self.milvus_client.connect()
                    logger.info("Milvus 重新连接成功")
                except Exception as e:
                    logger.error(f"Milvus 重新连接失败: {str(e)}")
            
            # 检查 PostgreSQL 可用性
            # 这里需要根据实际的 PostgreSQL 实现添加检查逻辑
            
            # 记录熔断器状态
            circuit_states = self.get_circuit_breaker_state()
            logger.debug(f"健康检查完成，熔断器状态: {circuit_states}")
            
        except Exception as e:
            logger.error(f"健康检查任务执行失败: {str(e)}")
    
    def _update_all_collection_sizes(self):
        """
        更新所有已知集合的大小
        """
        try:
            # 对于 Milvus，获取所有集合的统计信息
            if self.milvus_client._connected:
                for collection_name in list(self._collection_sizes.keys()):
                    try:
                        stats = self.milvus_client.get_collection_stats(collection_name)
                        if stats and 'row_count' in stats:
                            self.update_collection_size(collection_name, stats['row_count'])
                    except Exception as e:
                        logger.error(f"更新集合 {collection_name} 大小失败: {str(e)}")
            
            logger.debug("集合大小更新完成")
        except Exception as e:
            logger.error(f"更新集合大小任务执行失败: {str(e)}")


# 全局向量数据库路由器
vector_db_router = VectorDatabaseRouter()